import cv2
import numpy as np
import os
import traceback
from app.model.YoloModel import YoloModel

# === 参数配置 ===
OUTPUT = True
SAVE = False
model_path = '../models/yolov10n.pt'
video_path = '../videos/test.mp4'
calibration_file = '../app/calibration/lane_calibration.json'

# 事故检测参数
accident_detection = True
accident_types = ['off_road', 'collision']  # 指定要启用的检测器类型
accident_params = {
    'debug_enabled': True  # 启用调试模式
}

# 创建显示窗口
if OUTPUT:
    cv2.namedWindow("processed", cv2.WINDOW_NORMAL)
    cv2.namedWindow("row", cv2.WINDOW_NORMAL)
    cv2.namedWindow("birdView", cv2.WINDOW_NORMAL)

# === 模型初始化 ===
print("Initializing YOLO model...")
try:
    yoloModel = YoloModel(
        model_path=model_path,
        traffic_flow=True,
        num_lanes=2,
        max_trail_length=30,
        calibration_file=calibration_file,
        road_curvature=0.3,
        accident_detection=accident_detection,
        accident_params=accident_params,
        accident_types=accident_types  # 新增参数
    )

    if yoloModel.model is None:
        print("Error: Model failed to load!")
        exit(1)
    else:
        print(f"Model loaded successfully. Classes: {yoloModel.model.names}")
        print(f"事故检测类型: {accident_types}")

    print(f"固定鸟瞰图尺寸: {yoloModel.bird_view_width}x{yoloModel.bird_view_height}")

    if accident_detection:
        print("车祸检测功能已启用")
        print(f"启用的检测器: {len(yoloModel.accident_detectors)}")
    else:
        print("车祸检测功能已禁用")

except Exception as e:
    print(f"初始化YOLO模型失败: {str(e)}")
    traceback.print_exc()
    exit(1)

# === 打开视频流 ===
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
    print(f"Error: Could not open video {video_path}")
    exit(1)

fps = cap.get(cv2.CAP_PROP_FPS)
wait_time = int(1000 / fps) if fps > 0 else 30
print(f"Video FPS: {fps}, Wait time: {wait_time}ms")


# 工具函数
def resize_to_fit(img, max_width=800):
    if img is None or img.size == 0:
        return None
    h, w = img.shape[:2]
    if w > max_width:
        scale = max_width / w
        return cv2.resize(img, (int(w * scale), int(h * scale)))
    return img


# 轨迹显示辅助函数
def draw_trajectory_info(frame, active_tracks):
    cv2.putText(frame, f"Active Tracks: {active_tracks}",
                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)


# === 主处理循环 ===
frame_count = 0
while True:
    ret, frame = cap.read()
    frame_count += 1

    if not ret:
        print(f"End of video at frame {frame_count}")
        break

    if frame is None or frame.size == 0:
        print(f"Invalid frame at count {frame_count}")
        continue

    print(f"\nProcessing frame {frame_count} ({frame.shape[1]}x{frame.shape[0]})")

    # 初始化默认值
    active_tracks = 0
    processed_frame = frame.copy()
    row_frame = frame.copy()
    birdView_frame = np.zeros((yoloModel.bird_view_height, yoloModel.bird_view_width, 3), dtype=np.uint8)

    try:
        # 处理帧并获取跟踪结果
        processed_frame, row_frame, birdView_frame = yoloModel.track(frame)

        # 获取当前活跃的轨迹数量
        if hasattr(yoloModel, 'track_history'):
            active_tracks = len(yoloModel.track_history)
        else:
            active_tracks = 0

        # 在processed帧上添加轨迹计数信息
        draw_trajectory_info(processed_frame, active_tracks)

    except Exception as e:
        print(f"Tracking error: {str(e)}")
        traceback.print_exc()

        # 在processed帧上添加错误信息
        cv2.putText(processed_frame, "TRACKING ERROR", (50, 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        cv2.putText(processed_frame, str(e), (50, 100),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 创建固定尺寸的鸟瞰图
        yoloModel.draw_lane_lines(birdView_frame)

    # 显示
    if OUTPUT:
        disp_processed = resize_to_fit(processed_frame)
        disp_row = resize_to_fit(row_frame)
        disp_bird = resize_to_fit(birdView_frame)

        if disp_processed is not None:
            cv2.imshow("processed", disp_processed)
        if disp_row is not None:
            cv2.imshow("row", disp_row)
        if disp_bird is not None:
            # 在鸟瞰图底部显示累计事故数量和轨迹信息
            stats = yoloModel.get_statistics()
            accidents = stats.get('accident_details', [])

            # if accidents:
            #     cv2.putText(disp_bird, f"ACCIDENTS: {len(accidents)}",
            #                 (disp_bird.shape[1] - 300, 30),
            #                 cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

            # 添加轨迹计数信息
            cv2.putText(disp_bird, f"Tracks: {active_tracks}",
                        (10, disp_bird.shape[0] - 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

            cv2.imshow("birdView", disp_bird)

    # 键盘控制
    key = cv2.waitKey(wait_time) & 0xFF
    if key == ord('q'):
        stats = yoloModel.get_statistics()
        print("\n=== Final Statistics ===")
        print(f"Total frames processed: {frame_count}")
        print(f"Total vehicles detected: {stats.get('total_count', 0)}")
        print(f"Vehicle categories: {stats.get('category_count', {})}")
        print(f"Traffic flow crossings: {stats.get('crossing_count', 0)}")
        print(f"Currently tracked vehicles: {stats.get('current_tracks', 0)}")
        print(f"Total accidents detected: {stats.get('accident_count', 0)}")

        # 显示详细的事故历史
        accidents = stats.get('accident_details', [])
        if accidents:
            print("\nAccident History:")
            for i, accident in enumerate(accidents, 1):
                if isinstance(accident, tuple) and len(accident) >= 2:
                    print(f"  {i}. {accident[0]} - {accident[1]}")
                else:
                    print(f"  {i}. {accident}")
        break
    elif key == ord('r'):
        yoloModel.reset_statistics()
        print("Statistics reset")
    elif key == ord('a'):
        # 按'a'键显示事故历史
        stats = yoloModel.get_statistics()
        accidents = stats.get('accident_details', [])
        if accidents:
            print("\nCurrent Accident History:")
            for i, accident in enumerate(accidents, 1):
                if isinstance(accident, tuple) and len(accident) >= 2:
                    print(f"  {i}. {accident[0]} - {accident[1]}")
                else:
                    print(f"  {i}. {accident}")
        else:
            print("No accidents detected yet")
    elif key == ord('t'):
        # 显示当前跟踪信息
        if hasattr(yoloModel, 'track_history') and yoloModel.track_history:
            print("\nCurrent Track Information:")
            for track_id, track_data in yoloModel.track_history.items():
                if track_data:
                    last_position = track_data[-1] if isinstance(track_data[-1], (tuple, list)) and len(
                        track_data[-1]) >= 2 else (0, 0)
                    print(f"Track {track_id}: Position ({last_position[0]:.1f}, {last_position[1]:.1f})")
        else:
            print("No active tracks")
    # 新增键盘控制：调整阈值和最小帧数
    elif key == ord('+'):
        if accident_detection and hasattr(yoloModel, 'accident_detectors') and yoloModel.accident_detectors:
            for detector in yoloModel.accident_detectors:
                if hasattr(detector, 'threshold'):
                    detector.threshold += 5
                    print(f"边界阈值增加到: {detector.threshold}像素 (检测器: {type(detector).__name__})")
    elif key == ord('-'):
        if accident_detection and hasattr(yoloModel, 'accident_detectors') and yoloModel.accident_detectors:
            for detector in yoloModel.accident_detectors:
                if hasattr(detector, 'threshold'):
                    detector.threshold = max(0, detector.threshold - 5)
                    print(f"边界阈值减小到: {detector.threshold}像素 (检测器: {type(detector).__name__})")
    elif key == ord('f'):
        if accident_detection and hasattr(yoloModel, 'accident_detectors') and yoloModel.accident_detectors:
            for detector in yoloModel.accident_detectors:
                if hasattr(detector, 'min_out_frames'):
                    detector.min_out_frames += 1
                    print(f"最小连续帧数增加到: {detector.min_out_frames} (检测器: {type(detector).__name__})")
    elif key == ord('d'):
        if accident_detection and hasattr(yoloModel, 'accident_detectors') and yoloModel.accident_detectors:
            for detector in yoloModel.accident_detectors:
                if hasattr(detector, 'min_out_frames'):
                    detector.min_out_frames = max(1, detector.min_out_frames - 1)
                    print(f"最小连续帧数减小到: {detector.min_out_frames} (检测器: {type(detector).__name__})")
    elif key == ord('c'):
        # 切换事故检测类型
        if accident_types == ['off_road', 'collision']:
            accident_types = ['off_road']
            print("只启用离开道路检测")
        elif accident_types == ['off_road']:
            accident_types = ['collision']
            print("只启用碰撞检测")
        else:
            accident_types = ['off_road', 'collision']
            print("启用所有事故检测")

        # 重置模型以应用新配置
        yoloModel = YoloModel(
            model_path=model_path,
            traffic_flow=True,
            num_lanes=2,
            max_trail_length=30,
            calibration_file=calibration_file,
            road_curvature=0.3,
            accident_detection=accident_detection,
            accident_params=accident_params,
            accident_types=accident_types
        )

cap.release()
cv2.destroyAllWindows()
print("Program ended.")